/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

/**
 * @file greater.cpp
 */
#include "greater_tiling.h"
#include "register/op_def_registry.h"
#include "graph/utils/type_utils.h"
#include "tiling/platform/platform_ascendc.h"

namespace optiling
{
    const uint32_t BLOCK_SIZE = 32;
    constexpr uint32_t NUM = 16;
    constexpr int VAL_ZERO = 0;
    constexpr int X1_INDEX = 0;
    constexpr int X2_INDEX = 1;

    // broadcast params
    uint32_t reduce1[20] = {VAL_ZERO};
    uint32_t reduce2[20] = {VAL_ZERO};
    uint32_t shape[20] = {VAL_ZERO};
    uint32_t d = 1;
    uint32_t dim = VAL_ZERO;

    // tilling params
    uint32_t x1_length = VAL_ZERO;
    uint32_t x2_length = VAL_ZERO;
    uint32_t total_length = VAL_ZERO;
    uint32_t ALIGN_NUM = VAL_ZERO;
    uint32_t tiling_size = VAL_ZERO;
    uint32_t block_size = VAL_ZERO;
    uint32_t core_size = VAL_ZERO;
    uint32_t core_remain = VAL_ZERO;
    auto aivNum = 1;

    static void GetTensorShape(const gert::StorageShape *shape, std::vector<uint64_t> &inshapeVector, uint32_t outDimNum)
    {
        int n = outDimNum;
        for (int j = shape->GetStorageShape().GetDimNum() - 1; j >= 0; --j)
        {
            inshapeVector[--n] = shape->GetStorageShape().GetDim(j);
        }
    }

    static void GetBroadCastParams(std::vector<uint64_t> inshapeVector1, std::vector<uint64_t> inshapeVector2, gert::Shape outshape)
    {
        // Broadcast logic for two inputs
        for (int i = 0; i < outshape.GetDimNum(); i++)
        {
            shape[i] = outshape.GetDim(i);
            d *= outshape.GetDim(i);
            if (inshapeVector1[i] != outshape.GetDim(i)) reduce1[i] = 1;
            if (inshapeVector2[i] != outshape.GetDim(i)) reduce2[i] = 1;
        }

        dim = outshape.GetDimNum();
        for (int i = dim - 1; i >= 1; i--)
        {
            if (!reduce1[i - 1] && !reduce2[i - 1] && !reduce1[i] && !reduce2[i])
            {
                dim--;
                shape[i - 1] *= shape[i];
            }
            else
            {
                break;
            }
        }
        if (reduce1[dim - 1] || reduce2[dim - 1])
        {
            shape[dim] = 1;
            dim++;
        }
    }

    static void GetTillingParams(gert::TilingContext *context)
    {
        // 双输入版本调整
        uint32_t sizeofdatatype;
        uint64_t ub_size;
        auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
        auto socVersion = ascendcPlatform.GetSocVersion();
        ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);
        aivNum = ascendcPlatform.GetCoreNum();

        const int INPUT_TENSOR_COUNT = 2; // 修改为2个输入
        for (int i = 0; i < INPUT_TENSOR_COUNT; ++i)
        {
            total_length = std::max<uint32_t>(total_length, context->GetInputTensor(i)->GetShapeSize());
        }

        x1_length = context->GetInputShape(X1_INDEX)->GetStorageShape().GetShapeSize();
        x2_length = context->GetInputShape(X2_INDEX)->GetStorageShape().GetShapeSize();

        auto dt = context->GetInputDesc(0)->GetDataType();
        if (dt == ge::DT_FLOAT16 || dt == ge::DT_BF16)
        {
            sizeofdatatype = 2;
        }
        else if (dt == ge::DT_INT8 || dt == ge::DT_UINT8)
        {
            sizeofdatatype = 1;
        }
        else if (dt == ge::DT_INT64)
        {
            sizeofdatatype = 8;
        }
        else
        {
            sizeofdatatype = 4;
        }
        
        ALIGN_NUM = BLOCK_SIZE / sizeofdatatype;
        const int UB_DIV_COEF = 2;
        tiling_size = ((ub_size) / BLOCK_SIZE / UB_DIV_COEF) / NUM;
        tiling_size = tiling_size <= 64 ? tiling_size : tiling_size / 64 * 64;
        block_size = tiling_size * ALIGN_NUM;
        aivNum = (aivNum < total_length / block_size) ? aivNum : (total_length / block_size);
        aivNum = aivNum >= 1 ? aivNum : 1;
        if (aivNum == 0)
        {
            aivNum = 1;
        }
        core_size = (total_length / aivNum) / (ALIGN_NUM * 8) * (ALIGN_NUM * 8);
        core_remain = total_length - aivNum * core_size;


    }

    static ge::graphStatus TilingFunc(gert::TilingContext *context)
    {
        GreaterTilingData tiling;

        auto outshape = context->GetOutputShape(0)->GetOriginShape();
        uint32_t outDimNum = outshape.GetDimNum();
        const gert::StorageShape *shape1 = context->GetInputShape(0);
        const gert::StorageShape *shape2 = context->GetInputShape(1);
        std::vector<uint64_t> inshapeVector1(outDimNum, 1);
        std::vector<uint64_t> inshapeVector2(outDimNum, 1);

        GetTensorShape(shape1, inshapeVector1, outDimNum);
        GetTensorShape(shape2, inshapeVector2, outDimNum);

        bool need_broadcast = true;
        for (int i = 0; i < outshape.GetDimNum(); i++)
        {
            if (shape1->GetStorageShape().GetDim(i) != shape2->GetStorageShape().GetDim(i))
                need_broadcast = true;
        }
        if (need_broadcast)
        {
            GetBroadCastParams(inshapeVector1, inshapeVector2, outshape);
            tiling.set_shape(shape);
            tiling.set_reduce1(reduce1);
            tiling.set_reduce2(reduce2);
            tiling.set_dim(dim);
        }

        GetTillingParams(context);
        tiling.set_x1_length(x1_length);
        tiling.set_x2_length(x2_length);
        tiling.set_total_length(total_length);
        tiling.set_ALIGN_NUM(ALIGN_NUM);
        tiling.set_tiling_size(tiling_size);
        tiling.set_block_size(block_size);
        tiling.set_core_size(core_size);
        tiling.set_core_remain(core_remain);

        context->SetBlockDim(aivNum);
        context->SetTilingKey(need_broadcast ? 1 : 0);

        tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
        context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
        size_t *currentWorkspace = context->GetWorkspaceSizes(1);
        currentWorkspace[0] = 0;
        return ge::GRAPH_SUCCESS;
    }
}

namespace ge {
    static ge::graphStatus InferShape(gert::InferShapeContext* context)
    {
        const gert::Shape* x1_shape = context->GetInputShape(0);
        gert::Shape* y_shape = context->GetOutputShape(0);
        *y_shape = *x1_shape;
        return GRAPH_SUCCESS;
    }
    }
// namespace ge {
//     static ge::graphStatus InferShape(gert::InferShapeContext* context) {
//         const gert::Shape* x1_shape = context->GetInputShape(0);
//         const gert::Shape* x2_shape = context->GetInputShape(1);
//         gert::Shape* y_shape = context->GetOutputShape(0);
        
//         const uint32_t dim_num1 = x1_shape->GetDimNum();
//         const uint32_t dim_num2 = x2_shape->GetDimNum();
//         const uint32_t max_dim = std::max(dim_num1, dim_num2);
//         std::vector<int64_t> out_dims(max_dim, 1);
    
//         for (uint32_t i = 0; i < max_dim; ++i) {
//             const int64_t dim1 = (i < dim_num1) ? 
//                 x1_shape->GetDim(dim_num1 - max_dim + i) : 1;
//             const int64_t dim2 = (i < dim_num2) ? 
//                 x2_shape->GetDim(dim_num2 - max_dim + i) : 1;
    
//             // 新增广播规则检查
//             if (dim1 != dim2 && dim1 != 1 && dim2 != 1) {
//                 // VECTOR_INFER_SHAPE_INNER_ERR_REPORT(
//                 //     "Invalid broadcast at dim[%u], dim1=%ld, dim2=%ld", i, dim1, dim2);
//                 return GRAPH_FAILED;
//             }
//             out_dims[i] = (dim1 == 0 || dim2 == 0) ? 0 : std::max(dim1, dim2);
//         }
    
//         y_shape->SetDimNum(max_dim);
//         for (uint32_t i = 0; i < max_dim; ++i) {
//             y_shape->SetDim(i, out_dims[i]);
//         }
//         return GRAPH_SUCCESS;
//     }
//     } // namespace ge

namespace ops
{
    class Greater : public OpDef
    {
    public:
        explicit Greater(const char *name) : OpDef(name)
        {
            this->Input("x1")
                .ParamType(REQUIRED)
                .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT8, ge::DT_INT32, ge::DT_INT64, ge::DT_UINT8})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
            this->Input("x2")
                .ParamType(REQUIRED)
                .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT8, ge::DT_INT32, ge::DT_INT64, ge::DT_UINT8})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
            this->Output("y")
                .ParamType(REQUIRED)
                .DataType({ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});

            this->SetInferShape(ge::InferShape);

            this->AICore()
                .SetTiling(optiling::TilingFunc);
            this->AICore().AddConfig("ascend910b").AddConfig("ascend310p");
        }
    };

    OP_ADD(Greater);
}


// #include "greater_tiling.h"
// #include "register/op_def_registry.h"
// #include "graph/utils/type_utils.h"
// #include "tiling/platform/platform_ascendc.h"



// constexpr int32_t inputVarNum = 2;//输入变量
// const uint32_t BLOCK_SIZE = 32;
// const uint32_t BUFFER_NUM = 2;
// namespace optiling {
// static ge::graphStatus TilingFunc(gert::TilingContext* context)
// {
//     GreaterTilingData tiling;
//     // 广播相关参数计算
//     uint32_t x1_length = context->GetInputShape(0)->GetStorageShape().GetShapeSize();
//     uint32_t x2_length = context->GetInputShape(1)->GetStorageShape().GetShapeSize();
//     int64_t numshapes1 = context->GetInputShape(0)->GetStorageShape().GetDimNum();
//     int64_t numshapes2 = context->GetInputShape(1)->GetStorageShape().GetDimNum();
//     // 广播后维度数
//     int64_t maxShapeDim = std::max(numshapes1, numshapes2);
//     // 记录每个输入原始shape
//     int64_t shape[maxShapeDim * inputVarNum];
//     // 记录广播后shape
//     int64_t shapefull[maxShapeDim];
//      for (int k = 0; k < inputVarNum; ++k) {
//         int64_t *ss = &shape[k * maxShapeDim];
//         const gert::StorageShape* inputshape = context->GetInputShape(k);
//         int dim_num =  inputshape->GetStorageShape().GetDimNum();
//         for (int i = 0; i < maxShapeDim; i++) {
//             if(i < maxShapeDim - dim_num){
//                 ss[i] = 1;
//             }
//             else{
//             ss[i] = inputshape->GetStorageShape().GetDim(i - (maxShapeDim - dim_num));
//             }
//         }
//     }
//     uint32_t inputNum=1;
//     for (int k = 0; k < maxShapeDim; ++k) {
//         int64_t max_dim = 1;
//         int64_t *sf = &shapefull[0];
//         for (int i = 0; i < inputVarNum; ++i) {
//             int64_t *ss = &shape[i * maxShapeDim];
//             if (ss[k] > max_dim) {
//                 max_dim = ss[k];
//             }
//         }
//         sf[k] = max_dim;  // 广播后的形状
//         inputNum *= max_dim;
//     }

//     // 是否需要广播            0:不需要广播 1:需要广播
//     uint8_t is_broadcast = 1;
//     if (numshapes1 == numshapes2) {
//         bool all_dims_equal = true;
//         for (int64_t i = 0; i < numshapes1; ++i) {
//             if (context->GetInputShape(0)->GetStorageShape().GetDim(i) != context->GetInputShape(1)->GetStorageShape().GetDim(i)) {
//                 all_dims_equal = false;
//                 break;
//             }
//         }
//         is_broadcast = all_dims_equal ? 0 : 1;
//     } 

//     tiling.set_is_broadcast(is_broadcast);
//     tiling.set_maxShapeDim(maxShapeDim);
//     tiling.set_shape(shape);
//     tiling.set_shapefull(shapefull);
//     tiling.set_x1_length(x1_length);
//     tiling.set_x2_length(x2_length);
//     ///////////////////
//     uint64_t ubSize;
//     auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
//     ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
//     auto coreNum = ascendcPlatform.GetCoreNum();
    
//     // Based on the input length and the number of inputs, the number of bytes of the input data type is obtained
//     // uint32_t inputNum = context->GetInputShape(0)->GetStorageShape().GetShapeSize();
//     uint32_t typeLength = 0;
//     ge::TypeUtils::GetDataTypeLength(context->GetInputDesc(0)->GetDataType(), typeLength);
//     uint32_t inputLength = inputNum * typeLength;
//     uint32_t inputBytes = inputLength / inputNum;

//     // There are a total of 3 shared UB spaces in the input and output. If it's int8, there are 2 more TBUFs
//     uint32_t ubDataNumber = 16;
//     // The number of 32B data blocks that can be used for each data. DOUBLE BUFFER is already counted here
//     uint32_t tileBlockNum = (ubSize / BLOCK_SIZE / BUFFER_NUM) / ubDataNumber;
//     uint32_t tileDataNum = (tileBlockNum * BLOCK_SIZE) / inputBytes;

//     // Input data for 32B alignment
//     uint32_t inputLengthAlgin32 = (((inputLength + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE);
//     // There is at least 32B of data on each core, satisfying several settings for several cores. The maximum number of audits is the actual number of audits
//     coreNum = (coreNum <  inputLengthAlgin32 / BLOCK_SIZE) ? coreNum : inputLengthAlgin32 / BLOCK_SIZE;
//     coreNum = (coreNum >= 1) ? coreNum : 1;
//     uint32_t everyCoreInputBlockNum = inputLengthAlgin32 / BLOCK_SIZE / coreNum;
//     uint32_t tailBlockNum = (inputLengthAlgin32 / BLOCK_SIZE) % coreNum;
    
//     // Small chunks are calculated and sliced several times using the number of data on each core
//     uint32_t smallCoreDataNum = everyCoreInputBlockNum * BLOCK_SIZE / inputBytes;
//     uint32_t smallTileNum = everyCoreInputBlockNum / tileBlockNum;
//     uint32_t finalSmallTileNum = (everyCoreInputBlockNum % tileBlockNum) == 0 ? smallTileNum : smallTileNum + 1;
//     // Tail block calculation for small chunks of data
//     uint32_t smallTailDataNum = smallCoreDataNum - (tileDataNum * smallTileNum);
//     smallTailDataNum = smallTailDataNum == 0 ? tileDataNum : smallTailDataNum;
    
//     // The total length of a large block of data is 32B larger than that of a small block of data
//     everyCoreInputBlockNum += 1;
//     uint32_t bigCoreDataNum = everyCoreInputBlockNum * BLOCK_SIZE / inputBytes;
//     uint32_t bigTileNum = everyCoreInputBlockNum / tileBlockNum;
//     uint32_t finalBigTileNum = (everyCoreInputBlockNum % tileBlockNum) == 0 ? bigTileNum : bigTileNum + 1;
//     uint32_t bigTailDataNum = bigCoreDataNum - tileDataNum * bigTileNum;
//     bigTailDataNum = bigTailDataNum == 0 ? tileDataNum : bigTailDataNum;
    
//     tiling.set_smallCoreDataNum(smallCoreDataNum);
//     tiling.set_bigCoreDataNum(bigCoreDataNum);
//     tiling.set_tileDataNum(tileDataNum);
//     tiling.set_smallTailDataNum(smallTailDataNum);
//     tiling.set_bigTailDataNum(bigTailDataNum);
//     tiling.set_finalSmallTileNum(finalSmallTileNum);
//     tiling.set_finalBigTileNum(finalBigTileNum);
//     tiling.set_tailBlockNum(tailBlockNum);

//     context->SetBlockDim(coreNum);
//     tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
//     context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
//     return ge::GRAPH_SUCCESS;
// }
// }


// namespace ge {
// static ge::graphStatus InferShape(gert::InferShapeContext* context)
// {
//     const gert::Shape* x1_shape = context->GetInputShape(0);
//     gert::Shape* y_shape = context->GetOutputShape(0);
//     *y_shape = *x1_shape;
//     return GRAPH_SUCCESS;
// }
// }


// namespace ops {
// class Greater : public OpDef {
// public:
//     explicit Greater(const char* name) : OpDef(name)
//     {
//         this->Input("x1")
//             .ParamType(REQUIRED)
//             .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT8, ge::DT_INT32, ge::DT_INT64, ge::DT_UINT8})
//             .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
//             .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
//         this->Input("x2")
//             .ParamType(REQUIRED)
//             .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT8, ge::DT_INT32, ge::DT_INT64, ge::DT_UINT8})
//             .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
//             .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
//         this->Output("y")
//             .ParamType(REQUIRED)
//             .DataType({ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL})
//             .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
//             .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});

//         this->SetInferShape(ge::InferShape);

//         this->AICore()
//             .SetTiling(optiling::TilingFunc);
//         this->AICore().AddConfig("ascend910b").AddConfig("ascend310p");

//     }
// };

// OP_ADD(Greater);
// }
